from flask import Flask, request, jsonify
import base64
import cv2
import numpy as np
from ultralytics import YOLO
import os

# 全局模型变量
model = None

def init_model():
    """加载YOLO模型（只加载一次）"""
    global model
    # 使用原始路径进行模型加载
    model_path = r'D:\app\ohos\jiaotongbiaoshishibie\ultralytics-8.3.2\runs\train\exp6\weights\best.pt'
    print(f"正在加载模型: {model_path}")
    model = YOLO(model=model_path)
    print("模型加载完成!")
    return model

app = Flask(__name__)

def predict_image(image):
    """执行目标检测并返回结果"""
    try:
        # 执行目标检测
        results = model.predict(image, save=False, show=False)
        
        # 获取带标注的图像
        result_image = results[0].plot()  # 带检测框的图像
        
        # 将结果转换为Base64
        _, buffer = cv2.imencode('.jpg', result_image, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
        base64_result = base64.b64encode(buffer).decode('utf-8')
        
        # 解析检测结果
        detections = []
        if results[0].boxes is not None:
            for i, (box, conf, cls) in enumerate(zip(results[0].boxes.xyxy, results[0].boxes.conf, results[0].boxes.cls)):
                detections.append({
                    "id": i,
                    "class": results[0].names[int(cls.item())],
                    "confidence": round(conf.item(), 4),
                    "box": [round(c.item(), 1) for c in box]  # [x1, y1, x2, y2]
                })
        
        return {
            "base64_result": base64_result,
            "detections": detections,
            "success": True
        }
    except Exception as e:
        return {
            "error": f"处理失败: {str(e)}",
            "success": False
        }

@app.route('/detect', methods=['POST'])
def detect():
    """API端点：接收Base64图像，返回检测结果"""
    global model
    
    # 检查模型是否已加载
    if model is None:
        return jsonify({"error": "模型未初始化", "success": False}), 500
    
    # 获取请求数据
    data = request.get_json()
    if not data or 'image_base64' not in data:
        return jsonify({"error": "缺少image_base64参数", "success": False}), 400
    
    try:
        # 解码Base64图像
        base64_data = data['image_base64']
        if base64_data.startswith('data:image'):
            # 去除头部信息
            base64_data = base64_data.split(',', 1)[-1]
            
        image_bytes = base64.b64decode(base64_data)
        nparr = np.frombuffer(image_bytes, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        # 检测图像
        result = predict_image(image)
        
        # 返回结果
        return jsonify(result)
    
    except Exception as e:
        return jsonify({
            "error": f"请求处理失败: {str(e)}",
            "success": False
        }), 500

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    return jsonify({"status": "ok", "model_loaded": model is not None})

def run_server():
    """启动Flask服务"""
    # 绑定到所有网络接口（0.0.0.0）
    app.run(host='0.0.0.0', port=5000, threaded=True, debug=False)
    # 启动提示
    print(f"服务已启动: http://{get_local_ip()}:5000")

def get_local_ip():
    """获取本地IP地址"""
    try:
        import socket
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except:
        return "127.0.0.1"

if __name__ == '__main__':
    # 初始化模型
    model = init_model()
    
    # 启动API服务
    print("\nAPI服务启动成功!")
    print(f"本地访问地址: http://localhost:5000/health")
    print(f"网络访问地址: http://{get_local_ip()}:5000/health")
    print("检测接口: POST http://localhost:5000/detect")
    print("请求格式: {'image_base64': '图片的base64字符串'}\n")
    
    run_server()